/* Copyright © 2011 Anton Tolchanov <anton.tolchanov@gmail.com>
 *
 * Based on other gdnsd plugins by Brandon L Black <blblack@gmail.com>
 * and Jay Reitz <jreitz@gmail.com>
 *
 * This file is part of gdnsd-plugin-weighted.
 *
 * gdnsd-plugin-weighted is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * gdnsd-plugin-weighted is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with gdnsd.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

#include "config.h"
#define GDNSD_PLUGIN_NAME weighted
#include <gdnsd-plugin.h>

#include <sys/time.h>

// For pre-1.5.0 compat
#ifndef dmn_assert
#include <assert.h>
#define dmn_assert assert
#endif

#define MAX_ITEMS_PER_RESOURCE 256

typedef struct {
    uint32_t addr;
    monio_state_t state;
} addrstate_t;

typedef struct {
    addrstate_t* as;
    unsigned count;
} addrstates_t;

typedef struct {
    const char* name;
    uint8_t* cname;
    addrstates_t addrs;
    long weight;
} resource_item_t;

typedef struct {
    const char* name;
    bool is_cname;
    resource_item_t* items;
    unsigned count;
    long total_weight;
} resource_t;

static resource_t* resources = NULL;
static unsigned num_resources = 0;

static monio_list_t monio_list = { 0, NULL };

static bool bad_key(const char* key, unsigned klen, const vscf_data_t* d V_UNUSED, void* data) {
        log_fatal(data, key, klen);
}

monio_list_t* plugin_weighted_load_config(const vscf_data_t* config) {
    dmn_assert(config);
    dmn_assert(vscf_get_type(config) == VSCF_HASH_T);

    num_resources = vscf_hash_get_len(config);
    if (!num_resources)
        log_fatal("plugin_weighted: one or more resources must be defined");

    resources = malloc(num_resources * sizeof(resource_t));

    for (unsigned i = 0; i < num_resources; i++) {
        resource_t* this_res = &resources[i];
        unsigned res_name_len;
        this_res->name = strdup(vscf_hash_get_key_byindex(config, i, &res_name_len));
        const vscf_data_t* res_props = vscf_hash_get_data_byindex(config, i);
        if (vscf_get_type(res_props) != VSCF_HASH_T)
            log_fatal("plugin_weighted: the value of resource '%s' must be a hash", this_res->name);

        const vscf_data_t* addrs_data = vscf_hash_get_data_byconstkey(res_props, "addrs", true);
        const vscf_data_t* cnames_data = vscf_hash_get_data_byconstkey(res_props, "cnames", true);
        const vscf_data_t* addr_groups_data = vscf_hash_get_data_byconstkey(res_props, "addr_groups", true);

        unsigned num_values = 0;
        if ((addrs_data || addr_groups_data) && cnames_data)
            log_fatal("plugin_weighted: resource '%s': resource must be address or cname, not both", this_res->name);
        else if (addrs_data && addr_groups_data)
            log_fatal("plugin_weighted: resource '%s': resource should contain 'addrs' or 'addr_groups', not both", this_res->name);
        else if (!addrs_data && !cnames_data && !addr_groups_data)
            log_fatal("plugin_weighted: resource '%s': must contain 'addrs', 'addr_groups' or 'cnames' hash of datacenters to resources", this_res->name);
        else if (addrs_data && (VSCF_HASH_T != vscf_get_type(addrs_data) || !(num_values = vscf_hash_get_len(addrs_data))))
            log_fatal("plugin_weighted: resource '%s': 'addrs' must be defined as a hash of one or more datacenters to address mappings", this_res->name);
        else if (cnames_data && (VSCF_HASH_T != vscf_get_type(cnames_data) || !(num_values = vscf_hash_get_len(cnames_data))))
            log_fatal("plugin_weighted: resource '%s': 'cnames' must be defined as a hash of one or more datacenters to cname mappings", this_res->name);
        else if (addr_groups_data && (VSCF_HASH_T != vscf_get_type(addr_groups_data) || !(num_values = vscf_hash_get_len(addr_groups_data))))
            log_fatal("plugin_weighted: resource '%s': 'addr_groups' must be defined as a hash of one or more datacenters to address mappings", this_res->name);

        if (num_values > MAX_ITEMS_PER_RESOURCE)
            log_fatal("plugin_weighted: resource '%s': number of items cannot be more than %d", this_res->name, MAX_ITEMS_PER_RESOURCE);

        this_res->count = num_values;
        this_res->total_weight = 0;
        this_res->items = malloc(num_values * sizeof(resource_item_t));

        if (cnames_data) this_res->is_cname = true;
            else this_res->is_cname = false;

        const char* svctype = NULL;
        const vscf_data_t* res_stype = vscf_hash_get_data_byconstkey(res_props, "service_type", true);
        if (res_stype) {
            if (cnames_data) {
                log_fatal("plugin_weighted: resource '%s': service checks cannot be used for CNAME resources", this_res->name);
            } else {
                if (vscf_get_type(res_stype) != VSCF_SIMPLE_T)
                    log_fatal("plugin_weighted: resource '%s': 'service_type' must be a string if defined", this_res->name);
                svctype = vscf_simple_get_data(res_stype);
            }
        }

        for (unsigned j = 0; j < num_values; j++) {
            if (cnames_data) {
                const vscf_data_t* cname_data = vscf_hash_get_data_byindex(cnames_data, j);
                dmn_assert(cname_data);

                resource_item_t* resource_item = &this_res->items[j];
                resource_item->name = strdup(vscf_hash_get_key_byindex(cnames_data, j, NULL));

                if ( VSCF_ARRAY_T != vscf_get_type(cname_data)
                        || (2 != vscf_array_get_len(cname_data))
                        || VSCF_SIMPLE_T != vscf_get_type(vscf_array_get_data(cname_data, 0))
                        || VSCF_SIMPLE_T != vscf_get_type(vscf_array_get_data(cname_data, 1))
                        || !vscf_simple_get_as_long(vscf_array_get_data(cname_data, 1), &resource_item->weight))
                    log_fatal("plugin_weighted: resource '%s', item '%s': values must be arrays of [ CNAME string, WEIGHT int ]", this_res->name, resource_item->name);
                if (resource_item->weight < 0)
                    log_fatal("plugin_weighted: resource '%s', item '%s': weight cannot be negative", this_res->name, resource_item->name);

                uint8_t* dname = malloc(256);
                dname_status_t dnstat = vscf_simple_get_as_dname(vscf_array_get_data(cname_data, 0), dname);
                if (dnstat == DNAME_INVALID)
                    log_fatal("plugin_weighted: resource '%s', item '%s': CNAME is not a legal domainname", this_res->name, resource_item->name);
                if(dnstat == DNAME_VALID)
                    dname = dname_trim(dname);
                resource_item->cname = dname;
                this_res->total_weight+=resource_item->weight;
                log_debug("plugin_weighted: resource '%s', item '%s', CNAME '%s' added", this_res->name, resource_item->name, logf_dname(dname));
            } else if (addr_groups_data) {
                const vscf_data_t* addr_group_data = vscf_hash_get_data_byindex(addr_groups_data, j);
                dmn_assert(addr_group_data);

                resource_item_t* resource_item = &this_res->items[j];
                resource_item->name = strdup(vscf_hash_get_key_byindex(addr_groups_data, j, NULL));

                if ( VSCF_HASH_T != vscf_get_type(addr_group_data) )
                    log_fatal("plugin_weighted: resource '%s', item '%s': values must be hashes, containing 'weight' and 'addrs' keys", this_res->name, resource_item->name);

                const vscf_data_t* addr_group_weight_data = vscf_hash_get_data_byconstkey(addr_group_data, "weight", true);
                const vscf_data_t* addr_group_addrs_data = vscf_hash_get_data_byconstkey(addr_group_data, "addrs", true);

                if (!addr_group_weight_data)
                    log_fatal("plugin_weighted: resource '%s', item '%s': weight should be defined", this_res->name, resource_item->name);

                if ( VSCF_SIMPLE_T != vscf_get_type(addr_group_weight_data)
                        || !vscf_simple_get_as_long(addr_group_weight_data, &resource_item->weight))
                    log_fatal("plugin_weighted: resource '%s', item '%s': cannot parse weight", this_res->name, resource_item->name);

                if (resource_item->weight < 0)
                    log_fatal("plugin_weighted: resource '%s', item '%s': weight cannot be negative", this_res->name, resource_item->name);

                this_res->total_weight+=resource_item->weight;

                if (!addr_group_addrs_data)
                    log_fatal("plugin_weighted: resource '%s', item '%s': addrs should be defined", this_res->name, resource_item->name);

                if ( VSCF_HASH_T != vscf_get_type(addr_group_addrs_data))
                    log_fatal("plugin_weighted: resource '%s', item '%s': 'addrs' key should be a mapping of logical names to IPv4 addresses", this_res->name, resource_item->name);

                resource_item->addrs.count = vscf_hash_get_len(addr_group_addrs_data);

                if (!resource_item->addrs.count)
                    log_fatal("plugin_weighted: resource '%s', item '%s': 'addrs' should not be empty", this_res->name, resource_item->name);

                if (resource_item->weight % resource_item->addrs.count)
                    log_fatal("plugin_weighted: resource '%s', item '%s': weight of an item should be a power of %d (because the item contains %d address(es))", this_res->name, resource_item->name, resource_item->addrs.count, resource_item->addrs.count);

                resource_item->addrs.as = calloc(resource_item->addrs.count, sizeof(addrstate_t));

                for (unsigned k = 0; k < resource_item->addrs.count; k++) {
                    unsigned lb_name_len;
                    const char* lb_name = vscf_hash_get_key_byindex(addr_group_addrs_data, k, &lb_name_len);
                    const vscf_data_t* lb_data = vscf_hash_get_data_byindex(addr_group_addrs_data, k);
                    if (VSCF_SIMPLE_T != vscf_get_type(lb_data))
                        log_fatal("plugin_weighted: resource '%s', item '%s', name '%s': The values in the addrs hash must be IPv4 address strings", this_res->name, resource_item->name, lb_name);

                    const char* addr_txt = vscf_simple_get_data(lb_data);
                    struct in_addr a;
                    if(inet_pton(AF_INET, addr_txt, &a) < 1)
                        log_fatal("plugin_weighted: resource '%s', item '%s', name '%s': parsing '%s' as IPv4 address failed", this_res->name, resource_item->name, lb_name, addr_txt);   

                    resource_item->addrs.as[k].addr = a.s_addr;
                    log_debug("plugin_weighted: resource '%s', item '%s', address %s added", this_res->name, resource_item->name, addr_txt);

                    if (svctype) {
                        char *complete_desc = malloc(strlen(this_res->name) + strlen(resource_item->name) + lb_name_len + 3);
                        sprintf(complete_desc, "%s/%s/%s", this_res->name, resource_item->name, lb_name);
                        monio_list.info = realloc(monio_list.info, sizeof(monio_info_t) * (monio_list.count + 1));
                        monio_info_t* m = &monio_list.info[monio_list.count++];
                        m->svctype = svctype;
                        m->desc = complete_desc;
                        m->addr = addr_txt;
                        m->state_ptr = &resource_item->addrs.as[k].state;
                    } else {
                        // if no service type is defined, we just assign UP state to a resource
                        satom_set(&resource_item->addrs.as[0].state, MONIO_STATE_UP);
                    }

                }

                log_debug("plugin_weighted: resource '%s', item '%s' with %d addresses & weight %ld added", this_res->name, resource_item->name, resource_item->addrs.count, resource_item->weight);

            } else if (addrs_data) {
                const vscf_data_t* addr_data = vscf_hash_get_data_byindex(addrs_data, j);
                dmn_assert(addr_data);

                resource_item_t* resource_item = &this_res->items[j];
                resource_item->name = strdup(vscf_hash_get_key_byindex(addrs_data, j, NULL));

                if ( VSCF_ARRAY_T != vscf_get_type(addr_data)
                        || (2 != vscf_array_get_len(addr_data))
                        || VSCF_SIMPLE_T != vscf_get_type(vscf_array_get_data(addr_data, 0))
                        || VSCF_SIMPLE_T != vscf_get_type(vscf_array_get_data(addr_data, 1))
                        || !vscf_simple_get_as_long(vscf_array_get_data(addr_data, 1), &resource_item->weight))
                    log_fatal("plugin_weighted: resource '%s', item '%s': values must be arrays of [ IPADDR string, WEIGHT int ]", this_res->name, resource_item->name);
                if (resource_item->weight < 0)
                    log_fatal("plugin_weighted: resource '%s', item '%s': weight cannot be negative", this_res->name, resource_item->name);

                struct in_addr a;
                const char* addr_txt = vscf_simple_get_data(vscf_array_get_data(addr_data, 0));
                if (inet_pton(AF_INET, addr_txt, &a) < 1)
                    log_fatal("plugin_weighted: resource '%s', item '%s': parsing '%s' as IPv4 address failed", this_res->name, resource_item->name, addr_txt);
                resource_item->addrs.count = 1;
                resource_item->addrs.as = calloc(resource_item->addrs.count, sizeof(addrstate_t));
                resource_item->addrs.as[0].addr = a.s_addr;
                this_res->total_weight+=resource_item->weight;

                if (svctype) {
                    char *complete_desc = malloc(strlen(this_res->name) + strlen(resource_item->name) + 2);
                    sprintf(complete_desc, "%s/%s", this_res->name, resource_item->name);
                    monio_list.info = realloc(monio_list.info, sizeof(monio_info_t) * (monio_list.count + 1));
                    monio_info_t* m = &monio_list.info[monio_list.count++];
                    m->svctype = svctype;
                    m->desc = complete_desc;
                    m->addr = addr_txt;
                    m->state_ptr = &resource_item->addrs.as[0].state;
                } else {
                    // if no service type is defined, we just assign UP state to a resource
                    satom_set(&resource_item->addrs.as[0].state, MONIO_STATE_UP);
                }
                log_debug("plugin_weighted: resource '%s', item '%s': A '%s' added", this_res->name, resource_item->name, addr_txt);
            }
        }
        if (this_res->total_weight == 0)
            log_fatal("plugin_weighted: resource '%s': total weight of all items should not be 0", this_res->name);

        vscf_hash_iterate(res_props, true, bad_key, (void*)"plugin_weighted: Invalid resource option: '%s'");
    }

    return &monio_list;
}

unsigned plugin_weighted_map_resource_dyna(const char* resname) {
    if (!resname)
        log_fatal("plugin_weighted: resource name should be present");

    for (unsigned i = 0; i < num_resources; i++)
        if (!strcmp(resname, resources[i].name)) {
            if (resources[i].is_cname)
                log_fatal("plugin_weighted: Resource '%s' defined as CNAME data, but used in a DYNA RR", resources[i].name);
            log_debug("plugin_weighted: resource '%s' mapped", resources[i].name);
            return i;
        }

    log_fatal("plugin_weighted: unknown resource '%s'", resname);
}

unsigned plugin_weighted_map_resource_dync(const char* resname, const uint8_t* origin) {
    if (!resname)
        log_fatal("plugin_weighted: resource name should be present");
    for (unsigned i = 0; i < num_resources; i++) {
        if (!strcmp(resname, resources[i].name)) {
            if (!resources[i].is_cname)
                log_fatal("plugin_weighted: Resource '%s' defined as address data, but used in a DYNC RR", resources[i].name);
            for (unsigned j = 0; j < resources[i].count; j++) {
                const uint8_t* dname = resources[i].items[j].cname;
                if (dname_status(dname) == DNAME_PARTIAL) {
                    uint8_t dnbuf[256];
                    dname_copy(dnbuf, dname);
                    if (dname_cat(dnbuf, origin) != DNAME_VALID)
                        log_fatal("plugin_weighted: Name '%s' of resource '%s', when used at origin '%s', produces an invalid domainname", logf_dname(dname), resources[i].name, logf_dname(origin));
                }
            }
            log_debug("plugin_weighted: resource '%s' mapped", resources[i].name);
            return i;
        }
    }
    log_fatal("plugin_weighted: unknown resource '%s'", resname);
}

#if GDNSD_PLUGIN_API_VERSION > 5
void plugin_weighted_resolve_dyncname(unsigned threadnum V_UNUSED, unsigned resnum, const uint8_t* origin, const client_info_t* cinfo V_UNUSED, dyncname_result_t* result) {
#else
void plugin_weighted_resolve_dyncname(unsigned threadnum V_UNUSED, unsigned resnum, const uint8_t* origin, const anysin_t* client V_UNUSED, dyncname_result_t* result) {
#endif

    const resource_t* resource = &resources[resnum];
    dmn_assert(resource->is_cname);

    // We use the milliseconds value returned by gettimeofday as a pseudo-random number
    struct timeval tv;
    gettimeofday (&tv, 0);
    const unsigned long rand = (unsigned long)(tv.tv_usec) % resource->total_weight;

    unsigned long running_total = 0;
    for (unsigned j = 0; j < resource->count; j++) {
        running_total += resource->items[j].weight;
        if (running_total > rand) {
            const uint8_t* dname = resource->items[j].cname;
            dname_copy(result->dname, dname);
            if (dname_status(dname) == DNAME_PARTIAL) {
                dname_cat(result->dname, origin);
                dmn_assert(dname_status(result->dname) == DNAME_VALID);
            }
            return;
        }
    }

    // we should really never get here
    dmn_assert(0);
}

#if GDNSD_PLUGIN_API_VERSION > 5
void plugin_weighted_resolve_dynaddr(unsigned threadnum V_UNUSED, unsigned resnum, const client_info_t* cinfo V_UNUSED, dynaddr_result_t* result) {
#else
void plugin_weighted_resolve_dynaddr(unsigned threadnum V_UNUSED, unsigned resnum, const anysin_t* client V_UNUSED, dynaddr_result_t* result) {
#endif

    const resource_t* resource = &resources[resnum];
    dmn_assert(!resource->is_cname);
    result->count_v6 = 0;

    unsigned long good_item_count = 0; // number of good items
    unsigned long good_items[MAX_ITEMS_PER_RESOURCE];  // array of IDs of good items
    unsigned long good_items_weight[MAX_ITEMS_PER_RESOURCE]; // array of weights of good items
    unsigned long total_weight = 0;

    for (unsigned j = 0; j < resource->count; j++) {
        unsigned short int is_good = 0;
        for (unsigned i = 0; i < resource->items[j].addrs.count; i++) {
            if (monio_state_get(&resource->items[j].addrs.as[i].state) != MONIO_STATE_DOWN) {
                // Item is considered good if at least one IP address is not down
                if (is_good == 0) {
                    good_items[good_item_count] = j;
                    good_items_weight[good_item_count] = 0;
                    is_good = 1;
                }
                total_weight += (resource->items[j].weight / resource->items[j].addrs.count);
                good_items_weight[good_item_count] += (resource->items[j].weight / resource->items[j].addrs.count);
            }
        }
        if (is_good == 1) good_item_count++;
    }

    if (good_item_count == 0) {
        total_weight = resource->total_weight;
        good_item_count = resource->count;
        for (unsigned j = 0; j < resource->count; j++) {
            good_items[j] = j;
            good_items_weight[j] = resource->items[j].weight;
        }
    }

    struct timeval tv;
    gettimeofday (&tv, 0);
    const unsigned long rand = (unsigned long)(tv.tv_usec) % total_weight;

    unsigned long running_total = 0;
    for (unsigned j = 0; j<good_item_count; j++) {
        running_total += good_items_weight[j];
        if (running_total > rand) {
            for (unsigned i = 0; i < resource->items[good_items[j]].addrs.count; i++)
                if (monio_state_get(&resource->items[good_items[j]].addrs.as[i].state) != MONIO_STATE_DOWN)
                    result->addrs_v4[result->count_v4++] = resource->items[good_items[j]].addrs.as[i].addr;

            if (result->count_v4 == 0)
                // this means that all addresses of the item are in MONIO_STATE_DOWN, however the item is
                // considered good - this should happen only when all other items of the resource are down,
                // so we return IP addresses anyway
                for (unsigned i = 0; i < resource->items[good_items[j]].addrs.count; i++)
                    result->addrs_v4[result->count_v4++] = resource->items[good_items[j]].addrs.as[i].addr;

            return;
        }
    }

    // we should really never get here
    dmn_assert(0);
}

